import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

# ----------向上开口的曲线----------------
x = np.linspace(-5,5,20)
y = x ** 2 + 1

fig,axs = plt.subplots()

plt.ylim(-100,100)
plt.xlim(-10,10)
axs.plot(x,y,'r-')

#------------初始化的点以及切线------------
x_current = tf.Variable(-5.0)
y_current = x_current ** 2 + 1

# print(x_current.numpy())
# print(y_current.numpy())

line1, = axs.plot(x_current.numpy(), y_current.numpy(), 'bo')

x_tmp = x_current
with tf.GradientTape() as tape:
    y_tmp = x_tmp ** 2 + 1
    w = tape.gradient(y_tmp, x_tmp)
    b = y_tmp - w * x_tmp

    w = w.numpy()
    b = b.numpy()

    line2, = axs.plot(x, w * x + b, 'g--')
#------------
for i in range(10):
    # -----------变化线上的点的位置------------
    x_current = tf.Variable(-5.0 + i)
    y_current = x_current ** 2 + 1

    line1.set_data(x_current,y_current)

    # ----------变化点的切线---------------
    x_tmp = x_current
    with tf.GradientTape() as tape:
        y_tmp = x_tmp ** 2 + 1
        w = tape.gradient(y_tmp, x_tmp)
        b = y_tmp - w * x_tmp

        w = w.numpy()
        b = b.numpy()

        line2.set_data(x, w * x + b)

    # 每次间隔1s
    plt.pause(1)